from random import sample
from re import T
import os
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from sklearn.metrics import classification_report
from timm.data.mixup import Mixup
from timm.loss import SoftTargetCrossEntropy
from timm.models import swin_small_patch4_window7_224
import timm
from torchtoolbox.transform import Cutout
import albumentations
from albumentations.pytorch.transforms import ToTensorV2

# GPU ID 
GPU_ID = '1' # type : str
os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID

# experiment tag
extra_tag = './traffic_sign_ckpt/swin_small_patch4_window7_224'

# Global Parameter
ROOT_TRAIN_PATH = './data4train' # Dataset Path
model_lr = 1e-3  # Learning Rate
BATCH_SIZE = 32
EPOCHS = 100
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Classes num
num_of_cls = 10

# Data Augmentation and Preprocessing
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(), # p=0.5
    transforms.RandomVerticalFlip(), # p=0.5
    Cutout(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # imaegnet Normalization
])
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
mixup_fn = Mixup(
    mixup_alpha=0.8, cutmix_alpha=1.0, cutmix_minmax=None,
    prob=0.1, switch_prob=0.5, mode='batch',
    label_smoothing=0.1, num_classes=num_of_cls)

# Loading dataset
dataset_train = datasets.ImageFolder(ROOT_TRAIN_PATH +'/train', transform=transform)
dataset_test = datasets.ImageFolder(ROOT_TRAIN_PATH +'/val', transform=transform_test)
print(dataset_train.class_to_idx)

# Define DataLoader
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

#  Loading swin_transformer and Move to GPU
#  Loading ImageNet per-train model
criterion_train = SoftTargetCrossEntropy()
criterion_val = torch.nn.CrossEntropyLoss()
model_ft = swin_small_patch4_window7_224(pretrained=True) 
num_ftrs = model_ft.head.in_features
model_ft.head = nn.Linear(num_ftrs, num_of_cls)
model_ft.to(DEVICE)
print(model_ft)

# multi gpu setting
# if torch.cuda.device_count() > 1:
    # print("Let's use", torch.cuda.device_count(), "GPUs!")
    # model_ft = torch.nn.DataParallel(model_ft)

# optimizer setting
optimizer = optim.Adam(model_ft.parameters(), lr=model_lr)
cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=20, eta_min=1e-9)

# define train function
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    sum_loss = 0
    total_num = len(train_loader.dataset)
    print(f'Total trian num : {total_num}, iter_num: {len(train_loader)}')
    print('--------------------------------------------------------')
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
        samples, targets = mixup_fn(data, target)
        optimizer.zero_grad()
        output = model(samples)
        loss = criterion_train(output, targets)
        loss.backward()
        optimizer.step()
        lr = optimizer.state_dict()['param_groups'][0]['lr']
        print_loss = loss.data.item()
        sum_loss += print_loss
        # if (batch_idx + 1) % 10 == 0:
        print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR:{:.9f}'.format(
            epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                    100. * (batch_idx + 1) / len(train_loader), loss.item(), lr))
    ave_loss = sum_loss / len(train_loader)
    print('epoch:{},loss:{}'.format(epoch, ave_loss))



# define val function
ACC = 0
def val(model, device, test_loader):
    global ACC
    model.eval()
    test_loss = 0
    correct = 0
    total_num = len(test_loader.dataset)
    print(f'Total val num : {total_num}, iter_num: {len(test_loader)}')
    val_list = []
    pred_list = []
    with torch.no_grad():
        for data, target in test_loader:
            for t in target:
                val_list.append(t.data.item())
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion_val(output, target)
            _, pred = torch.max(output.data, 1)
            for p in pred:
                pred_list.append(p.data.item())
            correct += torch.sum(pred == target)
            print_loss = loss.data.item()
            test_loss += print_loss
        correct = correct.data.item()
        acc = correct / total_num
        avgloss = test_loss / len(test_loader)
        print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            avgloss, correct, len(test_loader.dataset), 100 * acc))
        if not os.path.exists(extra_tag):
            os.makedirs(extra_tag)
        torch.save(model_ft, f'./{extra_tag}/model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')
    return val_list, pred_list


# main
for epoch in range(1, EPOCHS + 1):
    train(model_ft, DEVICE, train_loader, optimizer, epoch)
    cosine_schedule.step()
    if epoch % 3 == 0 or epoch==EPOCHS: # per 3 epoch to validate
        val_list, pred_list = val(model_ft, DEVICE, test_loader)
        print(classification_report(val_list, pred_list, target_names=dataset_train.class_to_idx))